import os
import sys
import json
import math
import pickle
from typing import List, Dict, Tuple
import matplotlib
import matplotlib.pyplot as plt
from mpl_toolkits.mplot3d import Axes3D

# ensure repo root on path to import coord_transform if needed
SCRIPT_DIR = os.path.dirname(os.path.abspath(__file__))
REPO_ROOT = os.path.dirname(SCRIPT_DIR)
QUADCOPTER_DIR = os.path.join(REPO_ROOT, 'Quadcopter')
if QUADCOPTER_DIR not in sys.path:
    sys.path.insert(0, QUADCOPTER_DIR)

from coord_transform import PositionConvert
from SpaceAstarFixedWing_continuous import WorldMap, Kinematics, ContinuousAstarFixedWing

# Names that may indicate fixed-wing in template
FIXED_WING_TYPES = {'fixed_wing', 'fixed-wing', 'fixed_wing_uav', 'fixedwing'}

DEFAULT_VELOCITY = 20.0  # m/s fallback


def load_json(path: str) -> Dict:
    with open(path, 'r', encoding='utf-8') as f:
        return json.load(f)


def save_csv(uav_id: int, trajectory: List[Dict[str, float]], output_dir: str):
    os.makedirs(output_dir, exist_ok=True)
    filename = os.path.join(output_dir, f'UAV_{uav_id}_fixedwing_trajectory.csv')
    import csv
    with open(filename, 'w', newline='', encoding='utf-8') as cf:
        writer = csv.writer(cf)
        writer.writerow(['Latitude', 'Longitude', 'Altitude'])
        for p in trajectory:
            writer.writerow([p['lat'], p['lon'], p['alt']])
    print(f"Saved fixed-wing trajectory for UAV {uav_id} -> {filename}")
def plot_all_local_paths(all_paths: List[Tuple[int, List[Tuple[float,float,float]]]], output_dir: str):
    """
    Plot all UAV paths together in a single 3D figure.
    all_paths: list of (uav_id, path_xyz)
    """
    try:
        os.makedirs(output_dir, exist_ok=True)
        fig = plt.figure(figsize=(10, 8))
        ax = fig.add_subplot(111, projection='3d')
        cmap = plt.get_cmap('tab10')
        for i, (uav_id, path) in enumerate(all_paths):
            xs = [p[0] for p in path]
            ys = [p[1] for p in path]
            zs = [p[2] for p in path]
            color = cmap(i % 10)
            ax.plot(xs, ys, zs, lw=1.5, label=f'UAV{uav_id}', color=color)
            ax.scatter([xs[0]], [ys[0]], [zs[0]], marker='o', s=30, color=color)
            ax.scatter([xs[-1]], [ys[-1]], [zs[-1]], marker='^', s=40, color=color)
        ax.set_xlabel('X [m]')
        ax.set_ylabel('Y [m]')
        ax.set_zlabel('Z [m]')
        ax.set_title('All UAV Local Trajectories')
        ax.legend(loc='upper left', bbox_to_anchor=(1.02, 1))
        plt.tight_layout(rect=[0,0,0.85,1.0])
        outp = os.path.join(output_dir, 'all_uavs_local_trajectories.png')
        fig.savefig(outp, dpi=200, bbox_inches='tight')
        plt.show()
        print(f"Saved combined local plot -> {outp}")
    except ImportError:
        print('matplotlib not available, skipping combined plotting')
    except Exception as e:
        print(f'Failed to plot combined trajectories: {e}')


def parse_fixed_wing_tasks(template: Dict) -> Tuple[List[Dict], Dict]:
    tasks = []
    device_map = template.get('deviceModelParametersMap', {}) or {}
    # top environmental info (not used heavily here)
    env_top = template.get('environmentalInfo', {}) or {}

    lats = []
    lons = []
    alts = []

    for entry in template.get('routeToBePlannedInfos', []):
        dt = (entry.get('deviceType') or '').lower()
        if dt not in FIXED_WING_TYPES:
            continue
        device_model = entry.get('deviceModel')
        device_infos = entry.get('deviceInfos') or ([{'deviceId': entry.get('deviceId')}] if entry.get('deviceId') else [])
        if not device_infos:
            continue
        uav_id = device_infos[0].get('deviceId')

        sp = entry.get('startPoint')
        ep = entry.get('endPoint')
        wps = entry.get('waypoints') or entry.get('wayPoints') or []

        pts = []
        if sp:
            pts.append({'lat': sp.get('latitude'), 'lon': sp.get('longitude'), 'alt': sp.get('altitude', 0)})
            lats.append(sp.get('latitude'))
            lons.append(sp.get('longitude'))
            alts.append(sp.get('altitude', 0))
        for wp in wps:
            pts.append({'lat': wp.get('latitude'), 'lon': wp.get('longitude'), 'alt': wp.get('altitude', 0)})
            lats.append(wp.get('latitude'))
            lons.append(wp.get('longitude'))
            alts.append(wp.get('altitude', 0))
        if ep:
            pts.append({'lat': ep.get('latitude'), 'lon': ep.get('longitude'), 'alt': ep.get('altitude', 0)})
            lats.append(ep.get('latitude'))
            lons.append(ep.get('longitude'))
            alts.append(ep.get('altitude', 0))

        # device model velocity
        v_max = None
        if device_model and device_model in device_map:
            try:
                v_max = float(device_map[device_model].get('vMaxHorizon'))
            except Exception:
                v_max = None

        # time constraint similar handling
        extra = entry.get('extraInfo') or {}
        mft = None
        if extra.get('maxFlightTime') is not None:
            try:
                mft = float(extra.get('maxFlightTime'))
            except Exception:
                mft = None
        # build task
        task = {
            'uav_id': uav_id,
            'deviceModel': device_model,
            'v_max': v_max,
            'time_constraint': mft,
            'points': pts
        }
        tasks.append(task)

    bounds = None
    if lats and lons:
        bounds = {'min_lat': min(filter(None, lats)), 'max_lat': max(filter(None, lats)), 'min_lon': min(filter(None, lons)), 'max_lon': max(filter(None, lons)), 'min_alt': min(alts) if alts else 0, 'max_alt': max(alts) if alts else 100}
    return tasks, bounds


def build_fixedwing_worldmap(bounds: Dict, ref_lat: float, ref_lon: float, margin_m=500.0, resolution=1.0) -> WorldMap:
    # convert ref center to local coords: reference is ref_lat/ref_lon
    pc = PositionConvert()
    # convert bounds corners to XY
    min_x, min_y = pc.GPStoXY(bounds['min_lat'], bounds['min_lon'], ref_lat, ref_lon)
    max_x, max_y = pc.GPStoXY(bounds['max_lat'], bounds['max_lon'], ref_lat, ref_lon)
    # ensure min < max
    min_x, max_x = min(min_x, max_x), max(min_x, max_x)
    min_y, max_y = min(min_y, max_y), max(min_y, max_y)
    # apply margin
    min_x -= margin_m
    min_y -= margin_m
    max_x += margin_m
    max_y += margin_m
    # z bounds
    min_z = max(0.0, bounds.get('min_alt', 0) - 50)
    max_z = bounds.get('max_alt', 200) + 100

    world = WorldMap((min_x, min_y, min_z), (max_x, max_y, max_z), res=resolution, add_demo_obs=False)
    print(f"Built fixed-wing WorldMap: x [{min_x:.1f},{max_x:.1f}], y [{min_y:.1f},{max_y:.1f}], z [{min_z:.1f},{max_z:.1f}]")
    return world

def plan_fixedwing_multiseg(planner, pc, pts, ref_lat, ref_lon):
    """
    多段固定翼规划：
    输入 pts = [{'lat','lon','alt'}, ...]
    自动规划：
        pt0 -> pt1
        pt1 -> pt2
        ...
        pt(n-1) -> pt(n)

    返回: 全局路径 list[(x,y,z)]
    """
    full_path = []

    for i in range(len(pts) - 1):
        A = pts[i]
        B = pts[i + 1]

        # 转本地坐标
        Ax, Ay = pc.GPStoXY(A['lat'], A['lon'], ref_lat, ref_lon)
        Bx, By = pc.GPStoXY(B['lat'], B['lon'], ref_lat, ref_lon)
        Az = A['alt']
        Bz = B['alt']

        # 计算初始 yaw
        try:
            yaw = math.atan2(By - Ay, Bx - Ax)
            yaw_deg = math.degrees(yaw)
        except Exception:
            yaw_deg = 0.0
        # ===== 调用你的连续固定翼 A* 做规划 =====
        print(f"Segment {i}:{A['lat']:.6f},{A['lon']:.6f},{Az}  ->  {B['lat']:.6f},{B['lon']:.6f},{Bz}")
        print(f"x={Ax:.1f}, y={Ay:.1f}, z={Az:.1f}, yaw={yaw_deg:.1f}°  ->  x={Bx:.1f}, y={By:.1f}, z={Bz:.1f}")
        seg_path = planner.plan((Ax, Ay, Az, yaw_deg), (Bx, By, Bz), max_iter=200000)

        if seg_path is None:
            print(f"✘ Segment {i} FAILED")
            return None

        print(f"✓ Segment {i} OK, {len(seg_path)} points")

        # 拼接路径，避免重复点
        if i == 0:
            full_path.extend(seg_path)
        else:
            full_path.extend(seg_path[1:])

    return full_path

def plan_fixedwing_tasks(template_path: str, map_margin_m=500.0, resolution=1.0, output_dir='fixedwing_results', do_plot: bool = False):
    template = load_json(template_path)
    tasks, bounds = parse_fixed_wing_tasks(template)
    if not tasks:
        print('No fixed-wing tasks found in template')
        return

    # choose reference lat/lon as center of bounds
    ref_lat = 0.5 * (bounds['min_lat'] + bounds['max_lat'])
    ref_lon = 0.5 * (bounds['min_lon'] + bounds['max_lon'])
    pc = PositionConvert()

    # build simple worldmap for fixed-wing planning
    world = build_fixedwing_worldmap(bounds, ref_lat, ref_lon, margin_m=map_margin_m, resolution=resolution)

    # container for plotting combined results
    all_paths = []

    # plan each task
    for task in tasks:
        uav_id = task['uav_id']
        pts = task['points']
        if not pts or len(pts) < 2:
            print(f'UAV {uav_id} has insufficient points for planning')
            continue
        start = pts[0]
        goal = pts[-1]
        # compute yaw from start to next point
        next_point = pts[1] if len(pts) > 1 else goal
        # convert to local XY
        sx, sy = pc.GPStoXY(start['lat'], start['lon'], ref_lat, ref_lon)
        gx, gy = pc.GPStoXY(goal['lat'], goal['lon'], ref_lat, ref_lon)
        sz = start.get('alt', 0)
        gz = goal.get('alt', 0)
        dx = next_point['lat'] - start['lat']
        dy = next_point['lon'] - start['lon']
        yaw = 0.0
        try:
            # compute approximate yaw from local XY
            nx, ny = pc.GPStoXY(next_point['lat'], next_point['lon'], ref_lat, ref_lon)
            yaw = math.atan2(ny - sy, nx - sx)
            yaw_deg = math.degrees(yaw)
        except Exception:
            yaw_deg = 0.0

        # kinematics velocity from device model
        v = task.get('v_max') or DEFAULT_VELOCITY
        kin = Kinematics(max_climb_deg=15.0, velocity=v)

        planner = ContinuousAstarFixedWing(world, kin, step_len=2.0, yaw_discretization_deg=20.0,
                                           w=1.5, safe_dis=1.0, ignore_obstacles=True)

        print(f"Planning UAV {uav_id} start({start['lat']:.6f},{start['lon']:.6f},{sz}) -> goal({goal['lat']:.6f},{goal['lon']:.6f},{gz}), v={v}")
        print(f"Planning UAV {uav_id} with {len(pts)} waypoints...")

        # 多段规划（核心）
        full_path = plan_fixedwing_multiseg(planner, pc, pts, ref_lat, ref_lon)

        if full_path is None:
            print(f"UAV {uav_id} planning failed (some segment failed)")
            continue

        print(f"UAV {uav_id} multi-seg path OK, {len(full_path)} points")


        # convert path back to lat/lon/alt
        trajectory = []
        for x, y, z in full_path:
            lat, lon = pc.XYtoGPS(x, y, ref_lat, ref_lon)
            trajectory.append({'lat': lat, 'lon': lon, 'alt': z})

        save_csv(uav_id, trajectory, output_dir)

        if do_plot:
            all_paths.append((uav_id, full_path))

    # combined plot for all UAVs if requested
    if do_plot and all_paths:
        try:
            plot_all_local_paths(all_paths, output_dir)
        except Exception as e:
            print(f'Combined plotting failed: {e}')


if __name__ == '__main__':
    # default template.json in repo root
    tpl = os.path.join(REPO_ROOT, 'template.json')
    do_plot = True
    plan_fixedwing_tasks(tpl, do_plot=do_plot)
